/**
 * @file   MultigridSolver.cpp
 * @author HirasawaYui <yui@Ubuntu18-04>
 * @date   Tue Oct 27 21:38:49 2020
 * 
 * @brief  
 * 
 * 
 */

#include "MultigridSolver.h"

MultigridSolver::MultigridSolver(){
}

MultigridSolver::MultigridSolver(int n, double(*func)(double, double), double(*RS)(double, double), int Pre, int Post, double w)
{
    _n = n;
    _h = 1.0/pow(2, _n);
    _SdLen = (int)(pow(2, _n)) + 1;
    _func = func;
    _RealSltn = RS;
    for (int i = 0; i < _n; i++)
    {
	int length = (int)(pow(2, _n - i)) + 1;
	_f.push_back(std::vector<double>(length * length, 0));
	_v.push_back(std::vector<double>(length * length, 0));
    }
    _nowlevel = 1;
    _PreSmth = Pre;
    _PostSmth = Post;
    _w = w;
    _pRestrictOP = new FullWeightingRestriction();
    _pInterpolateOP = new LinearInterpolation();
}

int MultigridSolver::CortoIdx(int i, int j)
{
    return (i * _SdLen + j);
}

std::vector<int> MultigridSolver::IsBoundary()
{
    int size = _SdLen * _SdLen;
    std::vector<int> BndFlag(size, 0);
    for (int i = 0; i < _SdLen; i++)
    {
	BndFlag[i] = 1;
	BndFlag[size - 1 - i] = 1;
	BndFlag[i * _SdLen] = 1;
	BndFlag[size - i * _SdLen - 1] = 1;
    }
    return BndFlag;
}

void MultigridSolver::UpdateData()
{
    _SdLen = (int)(pow(2, _n - _nowlevel + 1)) + 1;
    _h = 1.0 / (_SdLen - 1);
    _BndMark = IsBoundary();
    _FEM = FEMSpace(_SdLen,_func,_RealSltn);
	_FEM.GenerateA();
	_FEM.GenerateRhs();
	_FEM.DealBoundary();
	
}

void MultigridSolver::WeightedJacobi(StiffMat A)
{
    std::vector<double> v_star(_SdLen*_SdLen, 0);
    std::map<int, double>::iterator it;
    for (int i = 0; i < _v[_nowlevel-1].size(); i++)
    {
	if (_BndMark[i] == 0)
	{
	    double a =_f[_nowlevel-1][i];
	    for (std::map<int, double>::iterator it = A[i].begin(); it != A[i].end(); it++)
	    {
		if (it->first == i)
		    continue;
		else
		{
		    a -= (it->second) * (_v[_nowlevel-1][it->first]);
		}
	    }
	    v_star[i] = a / A[i][i];
	}
	else
	    v_star[i] = _v[_nowlevel-1][i];
    }
    for (int i = 0; i < _v[_nowlevel-1].size(); i++)
	_v[_nowlevel-1][i] = (1.0 - _w) * _v[_nowlevel-1][i] + _w * v_star[i];	
}


void MultigridSolver::BottomSolve()
{
    _v[_nowlevel-1] = std::vector<double>(9, 0);
    _v[_nowlevel-1][4] = _f[_nowlevel-1][4] / _FEM.A()[4][4];
}

void MultigridSolver::VCycle(int StartLevel)
{
    _nowlevel = StartLevel;
    UpdateData();
    if (_nowlevel == _n)
    {
	BottomSolve();
    }
    else
    {
	for (int i = 0; i < _PreSmth; i++)
	    WeightedJacobi(_FEM.A());
	std::vector<double> r_h(_v[_nowlevel-1].size());
	std::vector<double> r_2h;
	// it;
	for (int i = 0; i < _v[_nowlevel-1].size(); i++)
	{
	    if (_BndMark[i] == 1)
		r_h[i] = 0;
	    else
	    {
		r_h[i] = _f[_nowlevel-1][i];
		for (std::map<int, double>::iterator it = _FEM._A[i].begin(); it != _FEM._A[i].end(); it++)
		    r_h[i] -= it->second * _v[_nowlevel-1][it->first];		
	    }
	}
	_pRestrictOP->SetInput(r_h);
	_pRestrictOP->SetSdLen(_SdLen);
	_pRestrictOP->restrict();
	r_2h = _pRestrictOP->ReturnOutput();
        _nowlevel++;
	UpdateData();
	_FEM.GenerateRhs();
	_FEM.DealBoundary();
	_f[_nowlevel-1] = r_2h;
	VCycle(_nowlevel);
	r_2h = _v[_nowlevel-1];
	_pInterpolateOP->SetInput(r_2h);
	_pInterpolateOP->SetSdLen(_SdLen);
	_pInterpolateOP->interpolate();
	r_h = _pInterpolateOP->ReturnOutput();
	_nowlevel--;
	UpdateData();
	for (int i = 0; i < _v[_nowlevel-1].size(); i++)
	    _v[_nowlevel-1][i] = _v[_nowlevel-1][i] + r_h[i];
	for (int i = 0; i < _PostSmth; i++)
	    WeightedJacobi(_FEM.A());	
    }
}
void MultigridSolver::FMG()
{
	UpdateData();
    if (_nowlevel == _n)
	{
		_v[_nowlevel-1] = std::vector<double>(9, 0);
		for(int i = 0;i < 9;i++)
		{
    		_v[_nowlevel-1][i] = _f[_nowlevel-1][i] / _FEM.A()[i][i];
		}
	} 
    else
    {
	std::vector<double> f_h = _f[_nowlevel-1];
	std::vector<double> f_2h;
	_pRestrictOP->SetInput(f_h);
	_pRestrictOP->SetSdLen(_SdLen);
	_pRestrictOP->restrict();
	f_2h = _pRestrictOP->ReturnOutput();
	_nowlevel++;
	UpdateData();
	_f[_nowlevel-1] = f_2h;
	FMG();
	f_2h = _v[_nowlevel-1];
	_pInterpolateOP->SetInput(f_2h);
	_pInterpolateOP->SetSdLen(_SdLen);
	_pInterpolateOP->interpolate();
	f_h = _pInterpolateOP->ReturnOutput();
	_nowlevel--;
	UpdateData();
	_v[_nowlevel-1] = f_h;
	for (int i = 0; i < _SdLen; i++)
	for (int j = 0; j < _SdLen; j++)
	{
	    int idx = CortoIdx(i, j);
	    if (_BndMark[idx] == 1)
		_v[_nowlevel-1][idx] = _RealSltn(i*_h, j*_h);
	}
	for (int i = 0; i < 1; i++)
	{
	    for (int j = _nowlevel; j < _n; j++)
	    {
		int length = (int)(pow(2, _n - j)) + 1;
		_f[j] = std::vector<double>(length * length, 0);
		_v[j] = std::vector<double>(length * length, 0);
	    }
	    VCycle(_nowlevel);
	}
    }
}
void MultigridSolver::Solve()
{
    _BndMark = IsBoundary();
    /// Set the initial value of _v and _f
    for (int i = 0; i < _SdLen; i++)
	for (int j = 0; j < _SdLen; j++)
	{
	    int idx = CortoIdx(i, j);
	    if (_BndMark[idx] == 1)
		_v[_nowlevel-1][idx] = _RealSltn(i*_h, j*_h);
	}
    _FEM = FEMSpace(_SdLen, _func,_RealSltn);
	_FEM.GenerateA();
    _FEM.GenerateRhs();
	_FEM.DealBoundary();
    _f[0] = _FEM.Rhs();
    for (int i = 0; i < 3; i++)
    {
        for (int j = 1; j < _n; j++)
		{
	    	int length = (int)(pow(2, _n - j)) + 1;
	    	_f[j] = std::vector<double>(length * length, 0);
	    	_v[j] = std::vector<double>(length * length, 0);
		}
	//VCycle(1);
	FMG();
	std::cout << "The iteration step: " << i+1 << ", the relative error: " << RE_maxnorm() << std::endl;
    }
    
}

double MultigridSolver::RE_2norm()
{
    double RSNorm = 0;
    double eNorm = 0;
    for (int i = 0; i < _SdLen; i++)
	for (int j = 0; j < _SdLen; j++)
	{
	    int idx = CortoIdx(i, j);
	    double a = _RealSltn(i*_h, j*_h);
	    RSNorm += pow(a, 2);
	    eNorm += pow(a - _v[0][idx], 2);
	}
    return sqrt(eNorm)/sqrt(RSNorm);
}

double MultigridSolver::RE_maxnorm()
{
	double RSNorm = 0;
    double eNorm = 0;
	for (int i = 0; i < _SdLen; i++)
	for (int j = 0; j < _SdLen; j++)
	{
	    int idx = CortoIdx(i, j);
	    double a = _RealSltn(i*_h, j*_h);
	    RSNorm = std::max(fabs(a),RSNorm);
	    eNorm = std::max(fabs(a - _v[0][idx]), eNorm);
	}
	return eNorm / RSNorm;
}
std::vector<double> MultigridSolver::ReturnSolution()
{
    return _v[0]; 
}

void MultigridSolver::WriteMatlabData()
{
    std::ofstream fs;
    fs.open("output.m");
    fs << "x = " << 0 << ":" << _h << ":" << 1 << ";" << std::endl;
    fs << "y = " << 0 << ":" << _h << ":" << 1 << ";" << std::endl;
    fs << "[X, Y] = meshgrid(x, y);" << std::endl;
    fs << "u = [";
    for (int i = 0; i < _SdLen; i++)
    {
        for (int j = 0; j < _SdLen; j++)
        {
            fs << _v[0][CortoIdx(i, j)] << " , ";
        }
        fs << ";" << std::endl;
    }
    fs << "];" << std::endl;
    fs << "surf(X, Y, u);" << std::endl;
}
